{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RK: this is outdated and does not work.\n",
    "#     Keep this ipynb as a placeholder for coreas_to_root.py\n",
    "#     Update this after coreas_to_root.py is complete.\n",
    "\n",
    "from __future__ import annotations\n",
    "\n",
    "from collections import OrderedDict\n",
    "from datetime import datetime\n",
    "from logging import getLogger\n",
    "import os\n",
    "from pathlib import Path\n",
    "import re\n",
    "from typing import Dict, Optional\n",
    "\n",
    "import astropy.constants\n",
    "from astropy.coordinates import CartesianRepresentation, PhysicsSphericalRepresentation\n",
    "from astropy.time import Time\n",
    "import astropy.units as u\n",
    "import numpy\n",
    "\n",
    "# from .generic import CollectionEntry, FieldsCollection, ShowerEvent\n",
    "from ipynb.fs.full.shower_generic import CollectionEntry, FieldsCollection, ShowerEvent\n",
    "\n",
    "# from ..antenna import ElectricField\n",
    "from ipynb.fs.full.antenna import ElectricField\n",
    "from grand.simulation.pdg import ParticleCode\n",
    "from grand.tools.coordinates import ECEF, LTP\n",
    "\n",
    "__all__ = [\"CoreasShower\"]\n",
    "\n",
    "\n",
    "logger = getLogger(__name__)\n",
    "\n",
    "\n",
    "\"\"\"CORSIKA particle Id to PDG code\"\"\"\n",
    "_id_to_code: Dict[int, ParticleCode] = {14: ParticleCode.PROTON, 5626: ParticleCode.IRON}\n",
    "\n",
    "\n",
    "class CoreasShower(ShowerEvent):\n",
    "    @classmethod\n",
    "    def _check_dir(cls, path: Path) -> bool:\n",
    "        try:\n",
    "            info_file = path.glob(\"*.reas\").__next__()\n",
    "        except StopIteration:\n",
    "            return False\n",
    "        return True\n",
    "\n",
    "    @classmethod\n",
    "    def _from_dir(cls, path: Path) -> CoreasShower:\n",
    "        if not path.exists():\n",
    "            raise FileNotFoundError(path)\n",
    "\n",
    "        matches = path.glob(\"*.reas\")\n",
    "        try:\n",
    "            reas_path = matches.__next__()\n",
    "        except StopIteration:\n",
    "            raise FileNotFoundError(path / \"*.reas\")\n",
    "        else:\n",
    "            index = int(reas_path.name[3:9])\n",
    "            reas = cls._parse_reas(path, index)\n",
    "\n",
    "            try:\n",
    "                matches.__next__()\n",
    "            except StopIteration:\n",
    "                pass\n",
    "            else:\n",
    "                logger.warning(f\"Multiple shower simulations in {path}. Loading only one.\")\n",
    "\n",
    "        config = {\n",
    "            \"energy\": float(reas[\"PrimaryParticleEnergy\"]) * 1e-09 << u.GeV,\n",
    "            \"zenith\": (180 - float(reas[\"ShowerZenithAngle\"])) << u.deg,\n",
    "            \"azimuth\": float(reas[\"ShowerAzimuthAngle\"]) << u.deg,\n",
    "            \"primary\": _id_to_code[int(reas[\"PrimaryParticleType\"])],\n",
    "        }\n",
    "\n",
    "        core = CartesianRepresentation(\n",
    "            float(reas[\"CoreCoordinateNorth\"]) * 1e-02,\n",
    "            float(reas[\"CoreCoordinateWest\"]) * 1e-02,\n",
    "            float(reas[\"CoreCoordinateVertical\"]) * 1e-02,\n",
    "            unit=u.m,\n",
    "        )\n",
    "        config[\"core\"] = core\n",
    "\n",
    "        geomagnet = PhysicsSphericalRepresentation(\n",
    "            theta=(90 + float(reas[\"MagneticFieldInclinationAngle\"])) << u.deg,\n",
    "            phi=0 << u.deg,\n",
    "            r=float(reas[\"MagneticFieldStrength\"]) * 1e-04 << u.T,\n",
    "        )\n",
    "        config[\"geomagnet\"] = geomagnet.represent_as(CartesianRepresentation)\n",
    "\n",
    "        distance = float(reas[\"DistanceOfShowerMaximum\"]) * 1e-02 << u.m\n",
    "        theta, phi = config[\"zenith\"], config[\"azimuth\"]\n",
    "        ct, st = numpy.cos(theta), numpy.sin(theta)\n",
    "        direction = CartesianRepresentation(st * numpy.cos(phi), st * numpy.sin(phi), ct)\n",
    "        config[\"maximum\"] = core - distance * direction\n",
    "\n",
    "        antpos = cls._parse_coreas_bins(path, index)\n",
    "        if antpos is None:\n",
    "            antpos = cls._parse_list(path, index)\n",
    "            if antpos is None:\n",
    "                antpos = cls._parse_info(path, index)\n",
    "\n",
    "        positions = {}\n",
    "        if antpos is not None:\n",
    "            for (antenna, r) in antpos:\n",
    "                positions[antenna] = CartesianRepresentation(x=r[0], y=r[1], z=r[2])\n",
    "\n",
    "        fields: Optional[FieldsCollection] = None\n",
    "        raw_fields = {}\n",
    "        try:\n",
    "            fields_path = path.glob(\"*_coreas\").__next__()\n",
    "        except StopIteration:\n",
    "            pass\n",
    "        else:\n",
    "            cgs2si = (astropy.constants.c / (u.m / u.s)).value * 1e02 * u.uV / u.m\n",
    "            pattern = re.compile(\"(\\d+).dat$\")\n",
    "            for antenna_path in fields_path.glob(\"*.dat\"):\n",
    "                antenna = int(pattern.search(str(antenna_path))[1])\n",
    "                logger.debug(f\"Loading trace for antenna {antenna}\")\n",
    "                data = numpy.loadtxt(antenna_path)\n",
    "                t = data[:, 0] * 1e09 * u.ns\n",
    "                Ex = data[:, 1] * cgs2si\n",
    "                Ey = data[:, 2] * cgs2si\n",
    "                Ez = data[:, 3] * cgs2si\n",
    "                electric = ElectricField(t, CartesianRepresentation(Ex, Ey, Ez), positions[antenna])\n",
    "                raw_fields[antenna] = CollectionEntry(electric)\n",
    "\n",
    "            fields = FieldsCollection()\n",
    "            for key in sorted(raw_fields.keys()):\n",
    "                fields[key] = raw_fields[key]\n",
    "\n",
    "        return cls(fields=fields, **config)\n",
    "\n",
    "    @classmethod\n",
    "    def _parse_reas(cls, path: Path, index: int) -> Optional[Dict]:\n",
    "        \"\"\"Parse a SIMxxxxxx.reas file\"\"\"\n",
    "        reas_file = path / f\"SIM{index:06d}.reas\"\n",
    "        if not reas_file.exists():\n",
    "            return None\n",
    "\n",
    "        with reas_file.open() as f:\n",
    "            txt = f.read()\n",
    "\n",
    "        try:\n",
    "            pattern = cls._reas_pattern\n",
    "        except AttributeError:\n",
    "            pattern = re.compile(\"([^=# \\n\\t]+)[ \\t]*=[ \\t]*([^ ;\\t]*)[ \\t]*;\")\n",
    "            setattr(cls, \"_reas_pattern\", pattern)\n",
    "\n",
    "        matches = pattern.findall(txt)\n",
    "\n",
    "        def tonum(s):\n",
    "            s2 = s[1:] if s.startswith(\"-\") else s\n",
    "            return int(s) if s2.isdecimal() else float(s)\n",
    "\n",
    "        return {key: tonum(value) for (key, value) in matches}\n",
    "\n",
    "    @classmethod\n",
    "    def _parse_coreas_bins(cls, path: Path, index: int) -> Optional[Dict]:\n",
    "        \"\"\"Parse a SIMxxxxxx_coreas.bins file\"\"\"\n",
    "        bins_file = path / f\"SIM{index:06d}_coreas.bins\"\n",
    "        if not bins_file.exists():\n",
    "            return None\n",
    "\n",
    "        data = []\n",
    "        pattern = re.compile(\"(\\d+).dat$\")\n",
    "        with bins_file.open() as f:\n",
    "            for line in f:\n",
    "                d = line.split()\n",
    "                if not d:\n",
    "                    continue\n",
    "                antenna = int(pattern.search(d[0])[1])\n",
    "                position = tuple(float(v) * 1e-02 for v in d[1:4]) << u.m\n",
    "                data.append((antenna, position))\n",
    "\n",
    "        return data\n",
    "\n",
    "    @classmethod\n",
    "    def _parse_list(cls, path: Path, index: int) -> Optional[Dict]:\n",
    "        \"\"\"Parse a SIMxxxxxx.list file\"\"\"\n",
    "\n",
    "        list_file = path / f\"SIM{index:06d}.list\"\n",
    "        if not list_file.exists():\n",
    "            return None\n",
    "\n",
    "        data = []\n",
    "        pattern = re.compile(\"(\\d+)$\")\n",
    "        with list_file.open() as f:\n",
    "            for line in f:\n",
    "                d = line.split()\n",
    "                if not d:\n",
    "                    continue\n",
    "                antenna = int(pattern.search(d[5])[1])\n",
    "                position = tuple(float(v) * 1e-02 for v in d[2:5]) << u.m\n",
    "                data.append((antenna, position))\n",
    "\n",
    "        return data\n",
    "\n",
    "    @classmethod\n",
    "    def _parse_info(cls, path: Path, index: int) -> Optional[Dict]:\n",
    "        \"\"\"Parse a SIMxxxxxx.info file\"\"\"\n",
    "\n",
    "        info_file = path / f\"SIM{index:06d}.info\"\n",
    "        if not info_file.exists():\n",
    "            return None\n",
    "\n",
    "        with info_file.open() as f:\n",
    "            txt = f.read()\n",
    "\n",
    "        try:\n",
    "            pattern = cls._info_pattern\n",
    "        except AttributeError:\n",
    "            pattern = re.compile(\n",
    "                \"ANTENNA[ \\t]+([^ \\t]+)[ \\t]+([^ \\t]+)\" \"[ \\t]+([^ \\t]+)[ \\t]+([^ \\t]+)\"\n",
    "            )\n",
    "            setattr(cls, \"_info_pattern\", pattern)\n",
    "\n",
    "        matches = pattern.findall(txt)\n",
    "\n",
    "        return [\n",
    "            (int(antenna), tuple(float(v) for v in values) << u.m) for (antenna, *values) in matches\n",
    "        ]"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Hide code",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
